{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "69ef446a-9d09-4016-b71d-a025b38a55d5",
   "metadata": {
    "tags": []
   },
   "source": [
    "## 向量数据库 embedding DB"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "82a74dfb-5e67-49db-829f-75214132793c",
   "metadata": {},
   "source": [
    "- 安装和部署\n",
    "- 数据库连接操作\n",
    "- 数据入库\n",
    "- 检索操作\n",
    "- 索引操作"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b91efcc1-5884-4d5d-b92c-fb22b3942d0d",
   "metadata": {
    "tags": []
   },
   "source": [
    "### 1.chroma\n",
    "- https://docs.trychroma.com/guides"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b0a1b4ac-5478-4e4d-a2a0-1eb87fe58d8a",
   "metadata": {},
   "source": [
    "- 安装和部署\n",
    "\n",
    "``` shell\n",
    "pip install chromadb\n",
    "```\n",
    "\n",
    "``` shell\n",
    "# 服务端部署\n",
    "chroma run --path ./data\n",
    "\n",
    " Usage: chroma run [OPTIONS]                                                                                                                        \n",
    " Run a chroma server                                                                                                                                \n",
    "╭─ Options ───────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╮\n",
    "│ --path            TEXT     The path to the file or directory. [default: ./chroma_data]                                                                                                                                                  │\n",
    "│ --host            TEXT     The host to listen to. Default: localhost [default: localhost]                                                                                                                                               │\n",
    "│ --log-path        TEXT     The path to the log file. [default: chroma.log]                                                                                                                                                              │\n",
    "│ --port            INTEGER  The port to run the server on. [default: 8000]                                                                                                                                                               │\n",
    "│ --help                     Show this message and exit.                                                                                                                                                                                  │\n",
    "╰─────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯\n",
    "\n",
    "```\n",
    "``` python\n",
    "客户端使用\n",
    "import chromadb\n",
    "chroma_client = chromadb.HttpClient(host='localhost', port=8000)\n",
    "```\n",
    "\n",
    "``` python\n",
    "# 直接使用\n",
    "import chromadb\n",
    "client = chromadb.Client()\n",
    "client = chromadb.PersistentClient(path=\"./data\")\n",
    "```\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "771d9427-c91d-4675-84b3-30732bc49732",
   "metadata": {},
   "outputs": [],
   "source": [
    "import chromadb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "3b0f9a41-d8ae-4b24-ab09-c75bd0f16fdb",
   "metadata": {},
   "outputs": [],
   "source": [
    "chroma_client = chromadb.HttpClient(host=\"localhost\", port=8000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "559f4a7a-c7d5-46c7-86ee-f9b02dfaef45",
   "metadata": {},
   "outputs": [],
   "source": [
    "from chromadb.utils import embedding_functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "50e35cba-fa42-40e3-bfee-8cf7316bc2a5",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_path = './data/llm_app/embedding_models/gte-large-zh/'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "46d5fdcf-9523-4e3d-b321-214df9f3a2cb",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/root/anaconda3/envs/llm/lib/python3.9/site-packages/sentence_transformers/cross_encoder/CrossEncoder.py:11: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from tqdm.autonotebook import tqdm, trange\n"
     ]
    }
   ],
   "source": [
    "em_fn = embedding_functions.SentenceTransformerEmbeddingFunction(model_name=model_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "2dcae39a-047c-43e8-81cb-c9f5f826f68b",
   "metadata": {},
   "outputs": [],
   "source": [
    "collection = chroma_client.create_collection(name='rag_db',\n",
    "                                            embedding_function=em_fn,\n",
    "                                            metadata={\"hnsw:space\": \"cosine\"})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "0ed9b0c6-9d8b-43d9-815c-e961e3ecbd24",
   "metadata": {},
   "outputs": [],
   "source": [
    "documents=[\"在向量搜索领域，我们拥有多种索引方法和向量处理技术，\\\n",
    "    它们使我们能够在召回率、响应时间和内存使用之间做出权衡。\", \n",
    "               \"虽然单独使用特定技术如倒排文件（IVF）、乘积量化（PQ）\\\n",
    "               或分层导航小世界（HNSW）通常能够带来满意的结果\",\n",
    "               \"GraphRAG 本质上就是 RAG，只不过与一般 RAG 相比，其检索路径上多了一个知识图谱\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "d03d17a8-1825-4d8c-ab00-6174b0ad9d8f",
   "metadata": {},
   "outputs": [],
   "source": [
    "collection.add(documents=documents,\n",
    "              ids=[\"id1\", \"id2\", \"id3\"],\n",
    "              metadatas=[{\"chapter\": 3, \"verse\": 16}, \n",
    "               {\"chapter\": 4, \"verse\": 5}, \n",
    "               {\"chapter\": 12, \"verse\": 5}])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "d5b41e8d-f416-4826-b57f-99f5af24fadc",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "3"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "collection.count()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "987cb02e-ac39-423e-a6c8-28d31356feb8",
   "metadata": {
    "collapsed": true,
    "jupyter": {
     "outputs_hidden": true
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'ids': ['id1'],\n",
       " 'embeddings': [[0.014445115812122822,\n",
       "   -0.011570136994123459,\n",
       "   0.022304952144622803,\n",
       "   0.02056988701224327,\n",
       "   -0.013929496519267559,\n",
       "   0.023787513375282288,\n",
       "   -0.0014493806520476937,\n",
       "   -0.022644491866230965,\n",
       "   0.004704768769443035,\n",
       "   -0.03234069421887398,\n",
       "   0.0016840423922985792,\n",
       "   -0.00019554801110643893,\n",
       "   0.01087995246052742,\n",
       "   0.031090212985873222,\n",
       "   0.005490060430020094,\n",
       "   0.010350426658987999,\n",
       "   -0.007863803766667843,\n",
       "   -0.002432334702461958,\n",
       "   -0.01740989089012146,\n",
       "   0.042129792273044586,\n",
       "   -0.020444413647055626,\n",
       "   0.02480866014957428,\n",
       "   0.04080279543995857,\n",
       "   -0.005154515150934458,\n",
       "   0.011707911267876625,\n",
       "   -0.007286903914064169,\n",
       "   0.0407462902367115,\n",
       "   0.002859817584976554,\n",
       "   0.010258081369102001,\n",
       "   -0.011978231370449066,\n",
       "   0.04582567512989044,\n",
       "   0.017384285107254982,\n",
       "   0.018167546018958092,\n",
       "   -0.029062245041131973,\n",
       "   0.059368252754211426,\n",
       "   -0.008747910149395466,\n",
       "   0.020769016817212105,\n",
       "   -0.025084445253014565,\n",
       "   0.023556025698781013,\n",
       "   0.0347982794046402,\n",
       "   0.01828758604824543,\n",
       "   0.02599630504846573,\n",
       "   -0.028328081592917442,\n",
       "   0.029358554631471634,\n",
       "   -0.033232640475034714,\n",
       "   -0.04206489771604538,\n",
       "   -0.001451569376513362,\n",
       "   0.04051527753472328,\n",
       "   -0.048479195684194565,\n",
       "   -0.012597007676959038,\n",
       "   0.026023663580417633,\n",
       "   -0.006464160047471523,\n",
       "   -0.01175509300082922,\n",
       "   0.017132123932242393,\n",
       "   -0.0008657105499878526,\n",
       "   -0.021280935034155846,\n",
       "   -0.025170451030135155,\n",
       "   -0.028863519430160522,\n",
       "   -0.01012085098773241,\n",
       "   0.013244359754025936,\n",
       "   -0.019075967371463776,\n",
       "   -0.09432715177536011,\n",
       "   -0.015523076988756657,\n",
       "   -0.0017137442482635379,\n",
       "   0.02409779466688633,\n",
       "   -0.02824191004037857,\n",
       "   -0.09161387383937836,\n",
       "   0.0039075445383787155,\n",
       "   -0.035326723009347916,\n",
       "   0.019154716283082962,\n",
       "   0.051437027752399445,\n",
       "   -0.018038827925920486,\n",
       "   -0.026541972532868385,\n",
       "   0.036077454686164856,\n",
       "   0.03859931230545044,\n",
       "   -0.01625019870698452,\n",
       "   -0.039681125432252884,\n",
       "   0.013105967082083225,\n",
       "   -0.015645213425159454,\n",
       "   -0.031582918018102646,\n",
       "   -0.04140772297978401,\n",
       "   -0.023032018914818764,\n",
       "   -0.012401299551129341,\n",
       "   -0.014670699834823608,\n",
       "   -0.008537613786756992,\n",
       "   -0.0089266961440444,\n",
       "   -0.012196388095617294,\n",
       "   -0.02505744807422161,\n",
       "   0.038365866988897324,\n",
       "   0.03193061053752899,\n",
       "   -0.03788961097598076,\n",
       "   0.03150743618607521,\n",
       "   -0.031922634690999985,\n",
       "   -0.02221149206161499,\n",
       "   0.008589781820774078,\n",
       "   -0.038415614515542984,\n",
       "   -0.05281243473291397,\n",
       "   0.04883595556020737,\n",
       "   0.03537963703274727,\n",
       "   0.01088634878396988,\n",
       "   -0.0563802495598793,\n",
       "   -0.04835273325443268,\n",
       "   0.01249020453542471,\n",
       "   0.013571920804679394,\n",
       "   0.017127668485045433,\n",
       "   -0.028376422822475433,\n",
       "   0.0005345133831724524,\n",
       "   -0.0013288226909935474,\n",
       "   -0.044484030455350876,\n",
       "   0.00855992455035448,\n",
       "   -0.016520436853170395,\n",
       "   0.04854075610637665,\n",
       "   0.01080259494483471,\n",
       "   0.004705204162746668,\n",
       "   -0.003448992967605591,\n",
       "   0.041245073080062866,\n",
       "   0.004165921360254288,\n",
       "   0.02815171517431736,\n",
       "   0.005700267851352692,\n",
       "   -0.011466828174889088,\n",
       "   0.03453798219561577,\n",
       "   0.0041418238542973995,\n",
       "   -0.025876102969050407,\n",
       "   -0.08596424758434296,\n",
       "   -0.018100585788488388,\n",
       "   0.05784747377038002,\n",
       "   -0.016291720792651176,\n",
       "   0.021245483309030533,\n",
       "   -0.02278871461749077,\n",
       "   -0.047689735889434814,\n",
       "   -0.0009158007451333106,\n",
       "   -0.01330314390361309,\n",
       "   -0.025066077709197998,\n",
       "   0.01401439867913723,\n",
       "   -0.0017432044260203838,\n",
       "   0.001487376051954925,\n",
       "   0.05186641961336136,\n",
       "   -0.018739258870482445,\n",
       "   -0.03384244814515114,\n",
       "   -0.03528480604290962,\n",
       "   -0.034695155918598175,\n",
       "   0.01063835620880127,\n",
       "   -0.019198233261704445,\n",
       "   -0.01410472672432661,\n",
       "   0.03006465919315815,\n",
       "   0.03892795741558075,\n",
       "   -0.0010317664127796888,\n",
       "   0.014635165221989155,\n",
       "   -0.055687129497528076,\n",
       "   -0.019645486027002335,\n",
       "   0.03855256363749504,\n",
       "   -0.0314452163875103,\n",
       "   -0.004341010935604572,\n",
       "   0.0061366441659629345,\n",
       "   -0.04147305339574814,\n",
       "   0.04228910058736801,\n",
       "   -0.13381339609622955,\n",
       "   0.006282264366745949,\n",
       "   0.003319332841783762,\n",
       "   -0.04294351115822792,\n",
       "   0.025187240913510323,\n",
       "   -0.026065319776535034,\n",
       "   0.039833828806877136,\n",
       "   0.006629409268498421,\n",
       "   0.017070012167096138,\n",
       "   0.02256811037659645,\n",
       "   0.09864982217550278,\n",
       "   -0.006643690634518862,\n",
       "   0.0030947932973504066,\n",
       "   0.003451511962339282,\n",
       "   -0.043554261326789856,\n",
       "   0.0034973756410181522,\n",
       "   0.016023680567741394,\n",
       "   0.03993913158774376,\n",
       "   -0.021862084046006203,\n",
       "   -0.005068321246653795,\n",
       "   -0.04628785327076912,\n",
       "   -0.013036703690886497,\n",
       "   -0.005984495859593153,\n",
       "   -0.016779806464910507,\n",
       "   -0.004195301327854395,\n",
       "   -0.002947553526610136,\n",
       "   0.023402446880936623,\n",
       "   0.05435528606176376,\n",
       "   -0.01838195137679577,\n",
       "   -0.02162923850119114,\n",
       "   -0.0024503206368535757,\n",
       "   0.03278221935033798,\n",
       "   0.014469532296061516,\n",
       "   -0.01953209564089775,\n",
       "   -0.01974594220519066,\n",
       "   0.009247761219739914,\n",
       "   -0.006697703618556261,\n",
       "   -0.03265515714883804,\n",
       "   -0.002428422449156642,\n",
       "   0.07099111378192902,\n",
       "   -0.03699566796422005,\n",
       "   0.0495075024664402,\n",
       "   0.018020762130618095,\n",
       "   0.010844447650015354,\n",
       "   -0.009493068791925907,\n",
       "   -0.011323750950396061,\n",
       "   0.045660845935344696,\n",
       "   0.0018337073270231485,\n",
       "   0.007658820599317551,\n",
       "   0.0332397036254406,\n",
       "   -0.0095193050801754,\n",
       "   0.0023129137698560953,\n",
       "   0.002042013220489025,\n",
       "   0.012419448234140873,\n",
       "   -0.016796601936221123,\n",
       "   -0.03163033351302147,\n",
       "   0.005510134156793356,\n",
       "   -0.016920235008001328,\n",
       "   -0.0009983867639675736,\n",
       "   0.07354562729597092,\n",
       "   0.03767923265695572,\n",
       "   0.015374156646430492,\n",
       "   -0.020589059218764305,\n",
       "   -0.01181066408753395,\n",
       "   -0.03817874565720558,\n",
       "   0.006068720947951078,\n",
       "   -0.009433514438569546,\n",
       "   -0.023960286751389503,\n",
       "   -0.0035128400195389986,\n",
       "   7.838066812837496e-05,\n",
       "   0.03718863055109978,\n",
       "   -0.04005344212055206,\n",
       "   -0.002586547052487731,\n",
       "   0.013273064978420734,\n",
       "   0.005569489672780037,\n",
       "   0.0049036904238164425,\n",
       "   -0.014088984578847885,\n",
       "   -0.06596098840236664,\n",
       "   0.02937222830951214,\n",
       "   0.02047060616314411,\n",
       "   0.01830219477415085,\n",
       "   0.003389948746189475,\n",
       "   -0.013638184405863285,\n",
       "   0.02579030580818653,\n",
       "   -0.00878108199685812,\n",
       "   -0.0009041570010595024,\n",
       "   0.004648271482437849,\n",
       "   -0.01905285008251667,\n",
       "   -0.01187395304441452,\n",
       "   0.0006751310429535806,\n",
       "   0.007919279858469963,\n",
       "   0.023353228345513344,\n",
       "   0.005151073448359966,\n",
       "   -0.018781600520014763,\n",
       "   -0.008139959536492825,\n",
       "   -0.09209931641817093,\n",
       "   0.054961368441581726,\n",
       "   0.05614292994141579,\n",
       "   -0.0030270421411842108,\n",
       "   0.048838526010513306,\n",
       "   -0.018354998901486397,\n",
       "   -0.008931172080338001,\n",
       "   -0.03626962751150131,\n",
       "   0.01860293187201023,\n",
       "   -0.023212764412164688,\n",
       "   0.05863054469227791,\n",
       "   -0.07010471075773239,\n",
       "   -0.002612892771139741,\n",
       "   0.04512390121817589,\n",
       "   -0.04634980484843254,\n",
       "   -0.02101266197860241,\n",
       "   0.004003269597887993,\n",
       "   -0.07885982096195221,\n",
       "   0.0945415273308754,\n",
       "   0.046234890818595886,\n",
       "   0.02023465745151043,\n",
       "   0.037477899342775345,\n",
       "   0.07870285212993622,\n",
       "   0.020105691626667976,\n",
       "   0.020908134058117867,\n",
       "   0.03668062016367912,\n",
       "   -0.08249036222696304,\n",
       "   0.00807509757578373,\n",
       "   0.02078310400247574,\n",
       "   0.027226926758885384,\n",
       "   -0.010630005970597267,\n",
       "   0.020788108929991722,\n",
       "   -0.013243497349321842,\n",
       "   0.029686562716960907,\n",
       "   0.05321459844708443,\n",
       "   -0.026125336065888405,\n",
       "   0.024448541924357414,\n",
       "   -0.007028741762042046,\n",
       "   0.04324077442288399,\n",
       "   0.0541263148188591,\n",
       "   -0.023722538724541664,\n",
       "   0.03270930051803589,\n",
       "   -0.011275696568191051,\n",
       "   -0.013297291472554207,\n",
       "   0.002133643953129649,\n",
       "   -0.018955683335661888,\n",
       "   -0.031040271744132042,\n",
       "   0.01016091275960207,\n",
       "   -0.004752648063004017,\n",
       "   -0.014330733567476273,\n",
       "   0.02414109744131565,\n",
       "   0.010987880639731884,\n",
       "   -0.003203909145668149,\n",
       "   0.03350331634283066,\n",
       "   0.007532360497862101,\n",
       "   0.039168328046798706,\n",
       "   -0.0068419454619288445,\n",
       "   0.025633731856942177,\n",
       "   0.020146220922470093,\n",
       "   -0.06663176417350769,\n",
       "   0.005901095923036337,\n",
       "   0.039277225732803345,\n",
       "   0.06177656725049019,\n",
       "   0.04843837395310402,\n",
       "   0.020706698298454285,\n",
       "   -0.03699494153261185,\n",
       "   -0.009119893424212933,\n",
       "   -0.03216291218996048,\n",
       "   -0.046463269740343094,\n",
       "   -0.0067017157562077045,\n",
       "   0.01056600734591484,\n",
       "   0.010992741212248802,\n",
       "   -0.004329986870288849,\n",
       "   -0.014998572878539562,\n",
       "   0.014632067643105984,\n",
       "   0.060682523995637894,\n",
       "   0.05631791800260544,\n",
       "   0.005588447209447622,\n",
       "   0.01121983677148819,\n",
       "   -0.06961537152528763,\n",
       "   -0.0043038493022322655,\n",
       "   -0.004681250546127558,\n",
       "   -0.02568148449063301,\n",
       "   -0.012509016320109367,\n",
       "   -0.03874756768345833,\n",
       "   0.042676739394664764,\n",
       "   -0.0911681056022644,\n",
       "   -0.0037112399004399776,\n",
       "   0.006496950518339872,\n",
       "   0.008282309398055077,\n",
       "   0.029424402862787247,\n",
       "   -0.039383988827466965,\n",
       "   -0.014376574195921421,\n",
       "   0.02548152394592762,\n",
       "   -0.004527352750301361,\n",
       "   -0.05091595649719238,\n",
       "   -0.04264511913061142,\n",
       "   0.00467529846355319,\n",
       "   0.010440748184919357,\n",
       "   0.020053347572684288,\n",
       "   0.03053177520632744,\n",
       "   -0.030955970287322998,\n",
       "   -0.006975631695240736,\n",
       "   -0.007110536098480225,\n",
       "   -0.02084576152265072,\n",
       "   0.038814641535282135,\n",
       "   -0.0010640432592481375,\n",
       "   0.04304073005914688,\n",
       "   0.025859659537672997,\n",
       "   0.025600876659154892,\n",
       "   -0.035846736282110214,\n",
       "   -0.0027339551597833633,\n",
       "   -0.032290469855070114,\n",
       "   0.04277687147259712,\n",
       "   0.0014725732617080212,\n",
       "   -0.06296177953481674,\n",
       "   0.05344432219862938,\n",
       "   0.002400320954620838,\n",
       "   -0.026633363217115402,\n",
       "   0.07473163306713104,\n",
       "   0.01095984224230051,\n",
       "   -0.0668288990855217,\n",
       "   0.0048528690822422504,\n",
       "   -0.047667670994997025,\n",
       "   -0.020861392840743065,\n",
       "   0.0018538878066465259,\n",
       "   -0.04565819352865219,\n",
       "   -0.024464815855026245,\n",
       "   0.02397610992193222,\n",
       "   -0.04716084897518158,\n",
       "   -0.026841171085834503,\n",
       "   -0.004968485329300165,\n",
       "   0.015656603500247,\n",
       "   -0.03032083809375763,\n",
       "   0.04403731971979141,\n",
       "   -0.04986147582530975,\n",
       "   -0.006502628792077303,\n",
       "   -0.0006518089794553816,\n",
       "   -0.03523901104927063,\n",
       "   -0.02697095088660717,\n",
       "   -0.06453543901443481,\n",
       "   -0.06082264333963394,\n",
       "   -0.026903821155428886,\n",
       "   0.009316083043813705,\n",
       "   0.045089490711688995,\n",
       "   0.024354659020900726,\n",
       "   -0.0018173549324274063,\n",
       "   0.04052475839853287,\n",
       "   -0.009512095712125301,\n",
       "   -0.020855385810136795,\n",
       "   -0.02947101555764675,\n",
       "   -0.026769572868943214,\n",
       "   0.014672295190393925,\n",
       "   0.024707255885004997,\n",
       "   -0.0009964160853996873,\n",
       "   -0.008160133846104145,\n",
       "   0.01818915829062462,\n",
       "   -0.024678630754351616,\n",
       "   0.030218198895454407,\n",
       "   0.011882291175425053,\n",
       "   -0.004989542067050934,\n",
       "   -0.06186346709728241,\n",
       "   -0.005053953267633915,\n",
       "   -0.054195601493120193,\n",
       "   -0.022457247599959373,\n",
       "   0.00664341589435935,\n",
       "   0.030270203948020935,\n",
       "   0.03298782557249069,\n",
       "   0.038533829152584076,\n",
       "   0.0031594401225447655,\n",
       "   -0.022814080119132996,\n",
       "   0.023959578946232796,\n",
       "   0.025023745372891426,\n",
       "   -0.030839595943689346,\n",
       "   -0.042113013565540314,\n",
       "   0.0771336480975151,\n",
       "   -0.016151167452335358,\n",
       "   -0.008188188076019287,\n",
       "   -0.014302707277238369,\n",
       "   0.04176957532763481,\n",
       "   -0.03794477880001068,\n",
       "   0.021580783650279045,\n",
       "   -0.02192535810172558,\n",
       "   -0.022079134359955788,\n",
       "   0.03231343254446983,\n",
       "   0.03848062828183174,\n",
       "   0.004814385902136564,\n",
       "   -0.025292987003922462,\n",
       "   0.002071705413982272,\n",
       "   -0.023004429414868355,\n",
       "   -0.030225081369280815,\n",
       "   -0.00497183995321393,\n",
       "   0.06423652172088623,\n",
       "   0.010873678140342236,\n",
       "   -0.021309170871973038,\n",
       "   0.012816035188734531,\n",
       "   -0.014305314049124718,\n",
       "   0.01280174683779478,\n",
       "   -0.04055511951446533,\n",
       "   -0.028355654329061508,\n",
       "   -0.034458406269550323,\n",
       "   0.010552824474871159,\n",
       "   -0.019365744665265083,\n",
       "   -0.01715274713933468,\n",
       "   -0.025457730516791344,\n",
       "   0.012138092890381813,\n",
       "   0.014913859777152538,\n",
       "   -0.034755729138851166,\n",
       "   0.012502927333116531,\n",
       "   0.051211144775152206,\n",
       "   0.037501584738492966,\n",
       "   0.020238889381289482,\n",
       "   -0.05019770935177803,\n",
       "   0.028107600286602974,\n",
       "   0.005764690227806568,\n",
       "   0.031999122351408005,\n",
       "   -0.04597993195056915,\n",
       "   -0.011727891862392426,\n",
       "   -0.022309767082333565,\n",
       "   0.01792989671230316,\n",
       "   -0.017555607482790947,\n",
       "   -0.01365154329687357,\n",
       "   0.03895464539527893,\n",
       "   -0.08405820280313492,\n",
       "   0.009110918268561363,\n",
       "   0.03814765810966492,\n",
       "   0.010010513477027416,\n",
       "   0.0479736253619194,\n",
       "   0.004327284172177315,\n",
       "   -0.016693009063601494,\n",
       "   -0.017989827319979668,\n",
       "   0.0031912866979837418,\n",
       "   0.013657419942319393,\n",
       "   0.023302428424358368,\n",
       "   0.03302807733416557,\n",
       "   -0.009803907945752144,\n",
       "   0.052340105175971985,\n",
       "   -0.002175351604819298,\n",
       "   0.012280349619686604,\n",
       "   0.01321385521441698,\n",
       "   -0.009757457301020622,\n",
       "   0.0008026933064684272,\n",
       "   -0.007939089089632034,\n",
       "   -0.007040894590318203,\n",
       "   -0.02380414493381977,\n",
       "   -0.04529658704996109,\n",
       "   -0.025621522217988968,\n",
       "   0.008006148040294647,\n",
       "   0.026646677404642105,\n",
       "   -0.04405933618545532,\n",
       "   -0.0019506871467456222,\n",
       "   -0.032955512404441833,\n",
       "   -0.015679821372032166,\n",
       "   0.025156527757644653,\n",
       "   0.03078193962574005,\n",
       "   0.04980805888772011,\n",
       "   0.03494549170136452,\n",
       "   -0.02120884880423546,\n",
       "   -0.013899730518460274,\n",
       "   -0.0169326514005661,\n",
       "   -0.05962421000003815,\n",
       "   0.028970859944820404,\n",
       "   0.023833956569433212,\n",
       "   0.02851969748735428,\n",
       "   -0.022233355790376663,\n",
       "   0.0372672937810421,\n",
       "   0.013318663462996483,\n",
       "   0.02460942044854164,\n",
       "   -0.021697577089071274,\n",
       "   -0.012667871080338955,\n",
       "   0.0981072261929512,\n",
       "   0.038865551352500916,\n",
       "   -0.03170158341526985,\n",
       "   0.06271585077047348,\n",
       "   -0.030487868934869766,\n",
       "   0.03493376821279526,\n",
       "   -0.0350528210401535,\n",
       "   -0.010324887931346893,\n",
       "   -0.02198502980172634,\n",
       "   0.02734844759106636,\n",
       "   0.03778421878814697,\n",
       "   -0.0223099235445261,\n",
       "   0.021038532257080078,\n",
       "   -0.0232114065438509,\n",
       "   0.028596945106983185,\n",
       "   0.006446299608796835,\n",
       "   -0.04709823429584503,\n",
       "   0.04552674666047096,\n",
       "   0.011034036986529827,\n",
       "   -0.06112547963857651,\n",
       "   0.015820862725377083,\n",
       "   -0.047209955751895905,\n",
       "   -0.03441028669476509,\n",
       "   -0.007523968815803528,\n",
       "   -0.048736877739429474,\n",
       "   0.01540645956993103,\n",
       "   0.03853219002485275,\n",
       "   0.010907004587352276,\n",
       "   -0.0012374905636534095,\n",
       "   0.04242873936891556,\n",
       "   0.0231510978192091,\n",
       "   0.030415650457143784,\n",
       "   0.03154049813747406,\n",
       "   0.04067664220929146,\n",
       "   -0.004507033620029688,\n",
       "   0.0008193549583666027,\n",
       "   0.07146511226892471,\n",
       "   0.047521136701107025,\n",
       "   -0.07420679926872253,\n",
       "   0.0017870953306555748,\n",
       "   0.0029683851171284914,\n",
       "   0.02042114920914173,\n",
       "   0.06229067221283913,\n",
       "   -0.02462659776210785,\n",
       "   -0.006266305223107338,\n",
       "   0.014872358180582523,\n",
       "   -0.011128094047307968,\n",
       "   0.025736162438988686,\n",
       "   -0.019007090479135513,\n",
       "   -0.017841657623648643,\n",
       "   -0.021756760776042938,\n",
       "   -0.09480136632919312,\n",
       "   -0.011473606340587139,\n",
       "   0.009662948548793793,\n",
       "   -0.012732379138469696,\n",
       "   -0.04481315612792969,\n",
       "   -0.020465917885303497,\n",
       "   -0.011116853915154934,\n",
       "   0.025861933827400208,\n",
       "   -0.01740095764398575,\n",
       "   0.051892880350351334,\n",
       "   0.02426142431795597,\n",
       "   -0.0670890137553215,\n",
       "   0.08187874406576157,\n",
       "   -0.0038591891061514616,\n",
       "   -0.010734659619629383,\n",
       "   0.031365346163511276,\n",
       "   -0.014165254309773445,\n",
       "   -0.08357790112495422,\n",
       "   0.02006210945546627,\n",
       "   0.0487162321805954,\n",
       "   -0.005874068476259708,\n",
       "   -0.00035432216827757657,\n",
       "   0.006092291325330734,\n",
       "   0.05460863560438156,\n",
       "   0.02514687180519104,\n",
       "   -0.017903290688991547,\n",
       "   0.014018689282238483,\n",
       "   0.025228649377822876,\n",
       "   -0.015697000548243523,\n",
       "   0.01454803068190813,\n",
       "   0.03505192697048187,\n",
       "   0.025768272578716278,\n",
       "   0.025690071284770966,\n",
       "   -0.003085083793848753,\n",
       "   -0.01798618584871292,\n",
       "   -0.03316478058695793,\n",
       "   0.0298068355768919,\n",
       "   0.04924459382891655,\n",
       "   0.03185010701417923,\n",
       "   -0.004788790829479694,\n",
       "   -0.0011440705275163054,\n",
       "   -0.013243059627711773,\n",
       "   -0.002620886079967022,\n",
       "   0.011527438648045063,\n",
       "   -0.031098760664463043,\n",
       "   -0.0670628622174263,\n",
       "   0.005088740028440952,\n",
       "   0.018483251333236694,\n",
       "   -0.017334194853901863,\n",
       "   0.053311627358198166,\n",
       "   0.005734212230890989,\n",
       "   0.04589887708425522,\n",
       "   0.03581713140010834,\n",
       "   0.05121186748147011,\n",
       "   0.025993140414357185,\n",
       "   -0.011597064323723316,\n",
       "   -0.025130748748779297,\n",
       "   0.0029494003392755985,\n",
       "   -0.0038464502431452274,\n",
       "   -0.01868605613708496,\n",
       "   0.027287857607007027,\n",
       "   0.039289262145757675,\n",
       "   -0.01966031827032566,\n",
       "   0.01632539927959442,\n",
       "   0.00497255427762866,\n",
       "   -0.0036921328864991665,\n",
       "   -0.004990049172192812,\n",
       "   -0.06674482673406601,\n",
       "   0.033885762095451355,\n",
       "   -0.0033720387145876884,\n",
       "   -0.04163259640336037,\n",
       "   -8.300475019495934e-05,\n",
       "   0.005929660052061081,\n",
       "   -0.0375823937356472,\n",
       "   -0.003595649264752865,\n",
       "   -0.0012355876388028264,\n",
       "   -0.007068119943141937,\n",
       "   0.007309237960726023,\n",
       "   0.06300930678844452,\n",
       "   -0.03452770784497261,\n",
       "   -0.01193853560835123,\n",
       "   0.022479653358459473,\n",
       "   0.004151070956140757,\n",
       "   -0.012026915326714516,\n",
       "   -0.027961833402514458,\n",
       "   -0.015191051177680492,\n",
       "   -0.01975247822701931,\n",
       "   0.013564814813435078,\n",
       "   -0.019076887518167496,\n",
       "   0.0652802586555481,\n",
       "   0.023616597056388855,\n",
       "   -0.033858396112918854,\n",
       "   0.017308669164776802,\n",
       "   0.04275878891348839,\n",
       "   -0.056962501257658005,\n",
       "   -0.08108912408351898,\n",
       "   -0.006717170588672161,\n",
       "   0.008761236444115639,\n",
       "   -0.01676054298877716,\n",
       "   -0.038817182183265686,\n",
       "   -0.002591116353869438,\n",
       "   -0.022149840369820595,\n",
       "   0.007824818603694439,\n",
       "   0.004599841311573982,\n",
       "   -0.006837856490164995,\n",
       "   -0.04252626374363899,\n",
       "   -0.02348034270107746,\n",
       "   -0.004914418328553438,\n",
       "   -0.012845797464251518,\n",
       "   -0.03573934733867645,\n",
       "   0.019550692290067673,\n",
       "   0.04157184436917305,\n",
       "   -0.004200706258416176,\n",
       "   -0.018032332882285118,\n",
       "   0.01368581410497427,\n",
       "   -0.02440798468887806,\n",
       "   0.019197458401322365,\n",
       "   0.014155609533190727,\n",
       "   -0.028440220281481743,\n",
       "   -0.012190978974103928,\n",
       "   -0.03715166449546814,\n",
       "   -0.02082945592701435,\n",
       "   0.10607072710990906,\n",
       "   0.012693522498011589,\n",
       "   0.04715524613857269,\n",
       "   0.03747210651636124,\n",
       "   -0.001302153104916215,\n",
       "   0.08100020885467529,\n",
       "   0.015899507328867912,\n",
       "   2.182554999308195e-05,\n",
       "   0.0033334344625473022,\n",
       "   0.023936735466122627,\n",
       "   0.02501005493104458,\n",
       "   -0.0034904060885310173,\n",
       "   -0.02775919996201992,\n",
       "   -0.06857065856456757,\n",
       "   -0.01928579993546009,\n",
       "   0.021284669637680054,\n",
       "   0.013801014982163906,\n",
       "   -0.028937706723809242,\n",
       "   -0.02280300296843052,\n",
       "   0.03438910096883774,\n",
       "   -0.017263047397136688,\n",
       "   0.02434469945728779,\n",
       "   0.007397370878607035,\n",
       "   -0.020997118204832077,\n",
       "   -0.050911445170640945,\n",
       "   0.02379179745912552,\n",
       "   -0.023917023092508316,\n",
       "   -0.02926446869969368,\n",
       "   0.003426693845540285,\n",
       "   0.03159898519515991,\n",
       "   0.029844745993614197,\n",
       "   -0.03458709269762039,\n",
       "   -0.031344711780548096,\n",
       "   0.05017240718007088,\n",
       "   0.046555615961551666,\n",
       "   -0.04456024244427681,\n",
       "   -0.04946606233716011,\n",
       "   -0.04809192568063736,\n",
       "   0.019160039722919464,\n",
       "   0.002831239951774478,\n",
       "   0.014653731137514114,\n",
       "   0.028612259775400162,\n",
       "   0.0033989783842116594,\n",
       "   -0.004297260195016861,\n",
       "   0.046990592032670975,\n",
       "   0.049215782433748245,\n",
       "   -0.012802842073142529,\n",
       "   0.015864796936511993,\n",
       "   -0.030538218095898628,\n",
       "   0.0075628128834068775,\n",
       "   -0.0037134240847080946,\n",
       "   0.01540209911763668,\n",
       "   0.050525765866041183,\n",
       "   -0.056665338575839996,\n",
       "   0.017494602128863335,\n",
       "   0.032926980406045914,\n",
       "   0.008067954331636429,\n",
       "   0.05018693953752518,\n",
       "   0.034181512892246246,\n",
       "   -0.057915642857551575,\n",
       "   0.04583194851875305,\n",
       "   0.003240117570385337,\n",
       "   0.03217657282948494,\n",
       "   0.0015755369095131755,\n",
       "   0.0021069813519716263,\n",
       "   0.009291416965425014,\n",
       "   0.0036583428736776114,\n",
       "   0.004013665020465851,\n",
       "   -0.009853566065430641,\n",
       "   -0.037380244582891464,\n",
       "   -0.03751889243721962,\n",
       "   -0.003744072513654828,\n",
       "   -0.002799970330670476,\n",
       "   -0.01937713846564293,\n",
       "   -0.0010906009702011943,\n",
       "   0.00011580262071220204,\n",
       "   -0.028204288333654404,\n",
       "   -0.03280114382505417,\n",
       "   0.0414755679666996,\n",
       "   0.010510805994272232,\n",
       "   0.0026559941470623016,\n",
       "   0.04423387348651886,\n",
       "   0.025830412283539772,\n",
       "   0.023717060685157776,\n",
       "   -0.011584130115807056,\n",
       "   0.04434794932603836,\n",
       "   -0.07935324311256409,\n",
       "   0.052767302840948105,\n",
       "   0.008063733577728271,\n",
       "   0.0051175267435610294,\n",
       "   -0.012480652891099453,\n",
       "   0.006169905420392752,\n",
       "   -0.0024866799358278513,\n",
       "   -0.008539539761841297,\n",
       "   -0.003386077471077442,\n",
       "   0.004373462405055761,\n",
       "   -0.011196954175829887,\n",
       "   -0.00828558299690485,\n",
       "   -0.012494356371462345,\n",
       "   -0.006706972140818834,\n",
       "   -0.0390702560544014,\n",
       "   -0.011064878664910793,\n",
       "   0.018109088763594627,\n",
       "   0.02439834363758564,\n",
       "   0.09074094146490097,\n",
       "   -0.004113936796784401,\n",
       "   -0.033617954701185226,\n",
       "   -0.013807570561766624,\n",
       "   0.00846595223993063,\n",
       "   0.02116408199071884,\n",
       "   -0.03550625964999199,\n",
       "   -0.04141256958246231,\n",
       "   0.029694579541683197,\n",
       "   -0.023130478337407112,\n",
       "   -0.031053487211465836,\n",
       "   -0.026472121477127075,\n",
       "   -0.021970465779304504,\n",
       "   -0.010224590077996254,\n",
       "   0.001577432849444449,\n",
       "   0.040241315960884094,\n",
       "   -0.024262161925435066,\n",
       "   -0.0018504903418943286,\n",
       "   -0.03306545689702034,\n",
       "   0.0275715421885252,\n",
       "   -0.0008220386225730181,\n",
       "   -0.008127396926283836,\n",
       "   -0.02106584422290325,\n",
       "   0.05293396860361099,\n",
       "   0.0200608279556036,\n",
       "   0.017029790207743645,\n",
       "   0.03306085616350174,\n",
       "   0.011605631560087204,\n",
       "   0.00868726335465908,\n",
       "   0.01677522249519825,\n",
       "   0.03515879064798355,\n",
       "   0.015194054692983627,\n",
       "   -0.020585007965564728,\n",
       "   -0.03946540132164955,\n",
       "   -0.0012483623577281833,\n",
       "   -0.013540997169911861,\n",
       "   0.04100408777594566,\n",
       "   -0.039479561150074005,\n",
       "   -0.006090169306844473,\n",
       "   0.023010052740573883,\n",
       "   0.0006918929284438491,\n",
       "   0.010698474012315273,\n",
       "   -0.016601400449872017,\n",
       "   0.04205101728439331,\n",
       "   0.021201176568865776,\n",
       "   -0.03984089195728302,\n",
       "   0.03355969488620758,\n",
       "   -0.01875278912484646,\n",
       "   -0.008728919550776482,\n",
       "   0.003694521961733699,\n",
       "   0.0013982439413666725,\n",
       "   -0.025383708998560905,\n",
       "   -0.05694817751646042,\n",
       "   -0.013089309446513653,\n",
       "   -0.008300725370645523,\n",
       "   -0.04031820967793465,\n",
       "   -0.01940823718905449,\n",
       "   -0.011046531610190868,\n",
       "   0.08344296365976334,\n",
       "   0.006322585977613926,\n",
       "   0.007630527019500732,\n",
       "   0.0001153216217062436,\n",
       "   0.03776954486966133,\n",
       "   0.036192476749420166,\n",
       "   0.042528536170721054,\n",
       "   0.02750822901725769,\n",
       "   0.007325693964958191,\n",
       "   -0.001694497070275247,\n",
       "   -0.024275463074445724,\n",
       "   -0.03922472149133682,\n",
       "   0.0201254952698946,\n",
       "   -0.004900798667222261,\n",
       "   -0.024712445214390755,\n",
       "   0.04312268644571304,\n",
       "   -0.012890060432255268,\n",
       "   -0.054043564945459366,\n",
       "   0.03597839176654816,\n",
       "   0.021126030012965202,\n",
       "   -0.03940432146191597,\n",
       "   -0.0804515928030014,\n",
       "   0.01826796680688858,\n",
       "   0.0008891065372154117,\n",
       "   0.05101679638028145,\n",
       "   -0.003752191783860326,\n",
       "   -0.00985755305737257,\n",
       "   -0.0047120461240410805,\n",
       "   0.00903234351426363,\n",
       "   -0.03332129493355751,\n",
       "   -0.014648513868451118,\n",
       "   0.023061614483594894,\n",
       "   -0.008648698218166828,\n",
       "   0.04963883012533188,\n",
       "   0.0008171244408003986,\n",
       "   0.0030519566498696804,\n",
       "   -0.05741933733224869,\n",
       "   0.01160670630633831,\n",
       "   -0.0062080323696136475,\n",
       "   -0.009905580431222916,\n",
       "   0.04484241083264351,\n",
       "   -0.012827744707465172,\n",
       "   -0.04862520098686218,\n",
       "   0.0312739759683609,\n",
       "   0.012228486128151417,\n",
       "   0.0010355050908401608,\n",
       "   0.00388712203130126,\n",
       "   -0.018810899928212166,\n",
       "   0.041598349809646606,\n",
       "   -0.04534130543470383,\n",
       "   0.01832222193479538,\n",
       "   -0.02372797578573227,\n",
       "   0.05154554918408394,\n",
       "   -0.04750785604119301,\n",
       "   0.022868676111102104,\n",
       "   -0.017851755023002625,\n",
       "   0.021791642531752586,\n",
       "   0.011872248724102974,\n",
       "   0.03867277503013611,\n",
       "   0.06091085076332092,\n",
       "   0.015602854080498219,\n",
       "   0.013781934045255184,\n",
       "   -0.027885470539331436,\n",
       "   0.038799528032541275,\n",
       "   -0.02039674110710621,\n",
       "   0.036257389932870865,\n",
       "   0.024567507207393646,\n",
       "   0.045997269451618195,\n",
       "   0.0075222039595246315,\n",
       "   -0.06575760245323181,\n",
       "   0.033040523529052734,\n",
       "   0.02895974926650524,\n",
       "   0.0008859183290041983,\n",
       "   0.06458812206983566,\n",
       "   0.007745934650301933,\n",
       "   0.001398889347910881,\n",
       "   0.019596876576542854,\n",
       "   -0.0072693065740168095,\n",
       "   0.008229678496718407,\n",
       "   -0.01914404146373272,\n",
       "   0.012830248102545738,\n",
       "   0.022337056696414948,\n",
       "   0.03082147054374218,\n",
       "   0.007175872102379799,\n",
       "   -0.012785721570253372,\n",
       "   0.042044829577207565,\n",
       "   -0.018223730847239494,\n",
       "   -0.028358934447169304,\n",
       "   -0.0005978968110866845,\n",
       "   0.017236867919564247,\n",
       "   -0.010528667829930782,\n",
       "   -0.039290621876716614,\n",
       "   -0.015173341147601604,\n",
       "   -0.039525825530290604,\n",
       "   -0.003886720398440957,\n",
       "   -0.028001144528388977,\n",
       "   0.03198913484811783,\n",
       "   0.013932991772890091,\n",
       "   0.03333419933915138,\n",
       "   -0.023178329691290855,\n",
       "   0.008042831905186176,\n",
       "   -0.0016310391947627068,\n",
       "   -0.011961312964558601,\n",
       "   -0.018025778234004974,\n",
       "   -0.08295862376689911,\n",
       "   -0.03249525651335716,\n",
       "   0.049558740109205246,\n",
       "   0.02339751087129116,\n",
       "   0.046556465327739716,\n",
       "   0.03638056665658951,\n",
       "   -0.024159733206033707,\n",
       "   0.001292865606956184,\n",
       "   -0.004171095322817564,\n",
       "   0.015539798885583878,\n",
       "   0.008430274203419685,\n",
       "   -0.0019029801478609443,\n",
       "   0.03378325328230858,\n",
       "   0.01749645173549652,\n",
       "   -0.01262248121201992,\n",
       "   0.03486904129385948,\n",
       "   -0.002285886090248823,\n",
       "   0.0038156635127961636,\n",
       "   0.006988596171140671,\n",
       "   -0.035656560212373734,\n",
       "   0.027769217267632484,\n",
       "   -0.03951457887887955,\n",
       "   -0.009207396768033504,\n",
       "   0.004153748042881489,\n",
       "   0.020918766036629677,\n",
       "   -0.031961582601070404,\n",
       "   0.022123554721474648,\n",
       "   0.007632121909409761,\n",
       "   -0.01798943802714348,\n",
       "   0.02556626684963703,\n",
       "   -0.03122851438820362,\n",
       "   -0.03735993802547455,\n",
       "   -0.036389973014593124,\n",
       "   0.0074442471377551556,\n",
       "   -0.014758375473320484,\n",
       "   -0.029871763661503792,\n",
       "   -0.0238412544131279,\n",
       "   0.023382600396871567,\n",
       "   0.020131923258304596,\n",
       "   -0.005771606229245663,\n",
       "   ...]],\n",
       " 'metadatas': [{'chapter': 3, 'verse': 16}],\n",
       " 'documents': ['在向量搜索领域，我们拥有多种索引方法和向量处理技术，    它们使我们能够在召回率、响应时间和内存使用之间做出权衡。'],\n",
       " 'data': None,\n",
       " 'uris': None,\n",
       " 'included': ['embeddings', 'documents', 'metadatas']}"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "collection.peek(limit=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f81e13ee-50cb-45be-8445-c2ad95e144c6",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0c8483e5-8de3-4f92-857e-b9b0c005400c",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "99799920-91aa-4ba6-9265-5c38318bbb03",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "7f5feedf-33bc-4d36-9d9d-b5343860503a",
   "metadata": {},
   "source": [
    "#### 检索"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "b19405b1-000d-4ba4-ba1c-a188eceeffde",
   "metadata": {},
   "outputs": [],
   "source": [
    "get_collection = chroma_client.get_collection(name='rag_db',\n",
    "                                             embedding_function=em_fn)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "37d03630-f355-4550-958a-cd170b8427ec",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "8667488b-79ef-41a8-a8c8-be2ca93ebadd",
   "metadata": {},
   "outputs": [],
   "source": [
    "id_result = get_collection.get(ids=['id2'],\n",
    "                              include=[\"documents\", \"embeddings\", \"metadatas\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "930995f1-aa44-44b5-a5f9-0f3d74472e62",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['虽然单独使用特定技术如倒排文件（IVF）、乘积量化（PQ）               或分层导航小世界（HNSW）通常能够带来满意的结果']"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "id_result['documents']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "99a80e5d-e430-4d9f-9403-fc0ed44a363f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[{'chapter': 4, 'verse': 5}]"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "id_result['metadatas']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "5d169446-8861-4d97-a4af-b1e68de57f72",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "294e662e-8832-4b93-a5cb-3a345c4bec25",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(1, 1024)"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.array(id_result['embeddings']).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "09daf84d-9418-48fa-a4ce-217fc56ae94d",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "f6976754-ff99-4fd0-be04-93933c916bf3",
   "metadata": {},
   "outputs": [],
   "source": [
    "query = '索引技术有哪些？'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "e12b5382-c326-4e16-943b-284457ad2c9a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'ids': [['id1', 'id3']],\n",
       " 'distances': None,\n",
       " 'embeddings': None,\n",
       " 'metadatas': [[{'chapter': 3, 'verse': 16}, {'chapter': 12, 'verse': 5}]],\n",
       " 'documents': [['在向量搜索领域，我们拥有多种索引方法和向量处理技术，    它们使我们能够在召回率、响应时间和内存使用之间做出权衡。',\n",
       "   'GraphRAG 本质上就是 RAG，只不过与一般 RAG 相比，其检索路径上多了一个知识图谱']],\n",
       " 'uris': None,\n",
       " 'data': None,\n",
       " 'included': ['documents', 'metadatas']}"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "get_collection.query(query_texts=query,\n",
    "                    n_results=2,\n",
    "                    include=[\"documents\", 'metadatas'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "51421098-4906-4b10-8bc5-3b6063ec17c9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'ids': [['id3', 'id2']],\n",
       " 'distances': None,\n",
       " 'embeddings': None,\n",
       " 'metadatas': [[{'chapter': 12, 'verse': 5}, {'chapter': 4, 'verse': 5}]],\n",
       " 'documents': [['GraphRAG 本质上就是 RAG，只不过与一般 RAG 相比，其检索路径上多了一个知识图谱',\n",
       "   '虽然单独使用特定技术如倒排文件（IVF）、乘积量化（PQ）               或分层导航小世界（HNSW）通常能够带来满意的结果']],\n",
       " 'uris': None,\n",
       " 'data': None,\n",
       " 'included': ['documents', 'metadatas']}"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "get_collection.query(query_texts=query,\n",
    "                    n_results=2,\n",
    "                    include=[\"documents\", 'metadatas'],\n",
    "                    where={\"verse\": 5})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "34c649f9-3b0b-496c-bc70-e548e798efe8",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "892b4696-c431-4eb8-bcb9-f7600a20c1f2",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "73037a4b-1d76-4917-a8ba-82b92fa6d8ab",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e1bb99fc-2854-4e43-ac25-675f967a6156",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "27549219-c12d-491a-a4de-c03d8c78995c",
   "metadata": {},
   "source": [
    "#### 混合检索支持的操作"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "626ab108-fc97-413b-bead-16faf7615030",
   "metadata": {},
   "source": [
    "\n",
    "\n",
    "``` shell\n",
    "- $eq - equal to (string, int, float)\n",
    "\n",
    "- $ne - not equal to (string, int, float)\n",
    "\n",
    "- $gt - greater than (int, float)\n",
    "\n",
    "- $gte - greater than or equal to (int, float)\n",
    "\n",
    "- $lt - less than (int, float)\n",
    "\n",
    "- $lte - less than or equal to (int, float)\n",
    "\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "dca17406-0bd0-4f00-b505-82a83b507ce8",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'ids': [['id1', 'id2']],\n",
       " 'distances': [[0.48257795562384265, 0.7083043108414732]],\n",
       " 'embeddings': None,\n",
       " 'metadatas': [[{'chapter': 3, 'verse': 16}, {'chapter': 4, 'verse': 5}]],\n",
       " 'documents': [['在向量搜索领域，我们拥有多种索引方法和向量处理技术，    它们使我们能够在召回率、响应时间和内存使用之间做出权衡。',\n",
       "   '虽然单独使用特定技术如倒排文件（IVF）、乘积量化（PQ）               或分层导航小世界（HNSW）通常能够带来满意的结果']],\n",
       " 'uris': None,\n",
       " 'data': None,\n",
       " 'included': ['metadatas', 'documents', 'distances']}"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "get_collection.query(\n",
    "    query_texts=[\"索引技术有哪些？\"],\n",
    "    n_results=2,\n",
    "    where={\"chapter\": {\"$lt\": 10}},\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "bd0ec614-ea9c-4cbd-9aee-292f5d3f80d2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'ids': [['id2']],\n",
       " 'distances': [[0.7083043108414732]],\n",
       " 'embeddings': None,\n",
       " 'metadatas': [[{'chapter': 4, 'verse': 5}]],\n",
       " 'documents': [['虽然单独使用特定技术如倒排文件（IVF）、乘积量化（PQ）               或分层导航小世界（HNSW）通常能够带来满意的结果']],\n",
       " 'uris': None,\n",
       " 'data': None,\n",
       " 'included': ['metadatas', 'documents', 'distances']}"
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "get_collection.query(\n",
    "    query_texts=[\"索引技术有哪些？\"],\n",
    "    n_results=2,\n",
    "    where={\"$and\": [{\"chapter\": {\"$lt\": 10}}, \n",
    "                    {\"verse\": {\"$eq\": 5}}\n",
    "                   ]}\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "3f4b70dc-29c4-4d6c-bb51-1382af0781c2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'ids': [['id1']],\n",
       " 'distances': [[0.48257795562384265]],\n",
       " 'embeddings': None,\n",
       " 'metadatas': [[{'chapter': 3, 'verse': 16}]],\n",
       " 'documents': [['在向量搜索领域，我们拥有多种索引方法和向量处理技术，    它们使我们能够在召回率、响应时间和内存使用之间做出权衡。']],\n",
       " 'uris': None,\n",
       " 'data': None,\n",
       " 'included': ['metadatas', 'documents', 'distances']}"
      ]
     },
     "execution_count": 33,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "get_collection.query(\n",
    "    query_texts=[\"索引技术有哪些？\"],\n",
    "    n_results=2,\n",
    "    where_document={\"$contains\":\"索引\"}\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "037c41b9-211e-4086-bfd7-fcdfef817f47",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "be7de6b6-e550-4dcf-a626-6dbcdf1f9870",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "b5e120b0-215c-42cc-b239-3e6281b788c6",
   "metadata": {
    "tags": []
   },
   "source": [
    "### 2.milvus\n",
    "- https://milvus.io/api-reference/pymilvus/v2.2.x/About.mdhttps://milvus.io/api-reference/pymilvus/v2.2.x/About.md"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d81a88d7-c136-4a91-8585-4e4a0c27a3c4",
   "metadata": {},
   "source": [
    "安装部署\n",
    "- https://milvus.io/docs/install-overview.md\n",
    "\n",
    "- docker 部署\n",
    "``` shell\n",
    "\n",
    "# 方式一\n",
    "curl -sfL https://raw.githubusercontent.com/milvus-io/milvus/master/scripts/standalone_embed.sh -o standalone_embed.sh\n",
    "\n",
    "bash standalone_embed.sh start\n",
    "bash standalone_embed.sh stop\n",
    "bash standalone_embed.sh delete\n",
    "\n",
    "\n",
    "```\n",
    "\n",
    "``` shell\n",
    "# 方式二\n",
    "# 通过docker-compose\n",
    "mkdir milvus_compose\n",
    "cd milvus_compose\n",
    "wget https://github.com/milvus-io/milvus/releases/download/v2.2.8/milvus-standalone-docker-compose.yml -O docker-compose.yml\n",
    " \n",
    "sudo systemctl daemon-reload\n",
    "sudo systemctl restart docker\n",
    "\n",
    "# 启动服务\n",
    "docker-compose up -d\n",
    "\n",
    "\n",
    "\n",
    "# 安装 python接口库\n",
    "pip install pymilvus\n",
    "\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "11b04691-5548-4e21-9927-8247693b3633",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from pymilvus import (\n",
    "    connections,\n",
    "    utility,\n",
    "    FieldSchema, CollectionSchema, DataType,\n",
    "    Collection,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "0abaef2b-fec9-4a51-a7db-21d91ec11ffc",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING: All log messages before absl::InitializeLog() is called are written to STDERR\n",
      "I0000 00:00:1728108761.480678   50952 config.cc:230] gRPC experiments enabled: call_status_override_on_cancellation, event_engine_dns, event_engine_listener, http2_stats_fix, monitoring_experiment, pick_first_new, trace_record_callops, work_serializer_clears_time_cache\n"
     ]
    }
   ],
   "source": [
    "connections.connect(host='127.0.0.1', port=\"19530\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "77d35caa-e518-4cc6-9eef-bfd4ba3640c9",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "fce0c849-82f7-428a-8e5e-da195dafa68b",
   "metadata": {},
   "source": [
    "#### 声明字段 和构建集合"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "d7d66184-3c8a-4a62-957a-3222017d1562",
   "metadata": {},
   "outputs": [],
   "source": [
    "fileds = [\n",
    "    FieldSchema(name=\"pk\", dtype=DataType.VARCHAR, \n",
    "                is_primary=True, auto_id=False, max_length=100),\n",
    "    FieldSchema(name=\"documents\", dtype=DataType.VARCHAR, max_length=512),\n",
    "    FieldSchema(name=\"embeddings\", dtype=DataType.FLOAT_VECTOR, dim=1024),\n",
    "    FieldSchema(name=\"verse\", dtype=DataType.INT64),\n",
    "    \n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "791e007e-76ca-4b2b-886b-4588746259a9",
   "metadata": {},
   "outputs": [],
   "source": [
    "rag_db = Collection(\"rag_db\",\n",
    "                    CollectionSchema(fileds),\n",
    "                   consistency_level=\"Strong\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6442a4b0-4574-486c-9063-b5385757422b",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ad9a6838-680b-4f93-8e59-c1821335ce2c",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "e322be46-9565-44e4-b0e6-66624f2f5952",
   "metadata": {},
   "source": [
    "#### 插入数据"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "6267878a-40ae-4c21-b285-fe7bf27dc63b",
   "metadata": {},
   "outputs": [],
   "source": [
    "documents=[\"在向量搜索领域，我们拥有多种索引方法和向量处理技术，\\\n",
    "    它们使我们能够在召回率、响应时间和内存使用之间做出权衡。\", \n",
    "               \"虽然单独使用特定技术如倒排文件（IVF）、乘积量化（PQ）\\\n",
    "               或分层导航小世界（HNSW）通常能够带来满意的结果\",\n",
    "               \"GraphRAG 本质上就是 RAG，只不过与一般 RAG 相比，其检索路径上多了一个知识图谱\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "id": "ef762c45-acd4-41b0-bb44-af49c4870c8b",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/tmp/ipykernel_50952/956235455.py:3: LangChainDeprecationWarning: The class `HuggingFaceEmbeddings` was deprecated in LangChain 0.2.2 and will be removed in 1.0. An updated version of the class exists in the langchain-huggingface package and should be used instead. To use it run `pip install -U langchain-huggingface` and import as `from langchain_huggingface import HuggingFaceEmbeddings`.\n",
      "  model = HuggingFaceEmbeddings(model_name=model_path,\n"
     ]
    }
   ],
   "source": [
    "from langchain.embeddings.huggingface import HuggingFaceEmbeddings\n",
    "model_path = './data/llm_app/embedding_models/gte-large-zh'\n",
    "model = HuggingFaceEmbeddings(model_name=model_path,\n",
    "                                   model_kwargs={'device': \"cpu\"})\n",
    "embeddings = model.embed_documents(documents)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "id": "168a838e-39b5-4d7d-9e4e-444b64d4829a",
   "metadata": {},
   "outputs": [],
   "source": [
    "entities = [\n",
    "    [str(i) for i in range(len(documents))],\n",
    "    documents,\n",
    "    np.array(embeddings),\n",
    "    [16,5,5],\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "id": "c90cf8bc-bbd6-4eee-863b-4fa1779e3e13",
   "metadata": {},
   "outputs": [],
   "source": [
    "insert_result = rag_db.insert(entities)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "id": "ca574d38-d038-46ed-a79b-ecccc79cd22d",
   "metadata": {},
   "outputs": [],
   "source": [
    "rag_db.flush()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "id": "cb4d0292-7c57-4245-b49c-929a97442a73",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "3"
      ]
     },
     "execution_count": 49,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "rag_db.num_entities"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "970cf628-ec35-4806-a1c1-7e37b95fb8bc",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "653dde8a-5aef-49cb-95fb-2616b507071d",
   "metadata": {},
   "source": [
    "#### 创建索引"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "id": "10041c5e-1e76-476b-b7b2-0ffd53b6b268",
   "metadata": {},
   "outputs": [],
   "source": [
    "index = {\n",
    "    \"index_type\": \"IVF_FLAT\",\n",
    "    \"metric_type\": \"L2\",\n",
    "    \"params\": {\"nlist\": 128},\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "id": "2496aa02-4b30-4d7b-a319-caa924589bdb",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Status(code=0, message=)"
      ]
     },
     "execution_count": 52,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "rag_db.create_index(\"embeddings\", index)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c6fad286-1913-46a7-8430-20e97ccacdda",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "01fb074a-b9b2-4267-b8fe-692cb391ea97",
   "metadata": {},
   "source": [
    "#### 检索"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "id": "323d70a5-9c1f-4331-9a07-e2f240cea68c",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "get_collection = Collection(\"rag_db\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "id": "749db02b-2e37-42a2-94fc-19429c94723a",
   "metadata": {},
   "outputs": [],
   "source": [
    "get_collection.load()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "id": "5eb83c12-04b8-4bdf-a1cd-650089d1f3fc",
   "metadata": {},
   "outputs": [],
   "source": [
    "query = \"索引技术有哪些？\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "id": "38b28261-5658-4f8b-89db-6d3a77db309d",
   "metadata": {},
   "outputs": [],
   "source": [
    "query_emb = model.embed_documents([query])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "id": "a604e489-366e-45e3-ba7d-78bb6081e6cb",
   "metadata": {},
   "outputs": [],
   "source": [
    "result = get_collection.search(query_emb, \n",
    "                       \"embeddings\", \n",
    "                       param={\"metric_type\": \"L2\"},\n",
    "                       limit=2, \n",
    "                       output_fields=[\"documents\", \"verse\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "id": "cb259648-9b67-4f44-a994-dafdfdf85ec5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "hit: (distance: 0.9651559591293335, id: 0), documents field: 在向量搜索领域，我们拥有多种索引方法和向量处理技术，    它们使我们能够在召回率、响应时间和内存使用之间做出权衡。\n",
      "hit: (distance: 1.1978662014007568, id: 2), documents field: GraphRAG 本质上就是 RAG，只不过与一般 RAG 相比，其检索路径上多了一个知识图谱\n"
     ]
    }
   ],
   "source": [
    "for hits in result:\n",
    "    for hit in hits:\n",
    "        print(f\"hit: {hit}, documents field: {hit.entity.get('documents')}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "id": "30c5b6b2-755b-4741-941a-a80fb12250c0",
   "metadata": {},
   "outputs": [],
   "source": [
    "result2 = get_collection.search(query_emb, \n",
    "                       \"embeddings\", \n",
    "                       param={\"metric_type\": \"L2\"},\n",
    "                       expr=\"verse < 10\",\n",
    "                       limit=1, \n",
    "                       output_fields=[\"documents\", \"verse\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "id": "5e1538c0-f54b-49b4-9e5e-188c9f28dc3f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "hit: (distance: 1.1978662014007568, id: 2), documents field: GraphRAG 本质上就是 RAG，只不过与一般 RAG 相比，其检索路径上多了一个知识图谱\n"
     ]
    }
   ],
   "source": [
    "for hits in result2:\n",
    "    for hit in hits:\n",
    "        print(f\"hit: {hit}, documents field: {hit.entity.get('documents')}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9017af5a-20cc-442f-a9eb-8c00dfffd4ad",
   "metadata": {
    "tags": []
   },
   "source": [
    "### 3.INDEX索引优化"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "id": "00f7e7bb-24ee-49bc-ad18-8b625c3cffcf",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from scipy.cluster.vq import kmeans2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "id": "05a8c595-508a-41ea-b5e0-19df7d4e1554",
   "metadata": {},
   "outputs": [],
   "source": [
    "query = np.random.normal(size=(128,))\n",
    "dataset = np.random.normal(size=(1000, 128))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d10627cd-1bf5-476c-961f-53a651ea70fe",
   "metadata": {
    "tags": []
   },
   "source": [
    "#### flat index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "id": "62f96c0e-67a4-472c-b73a-ddf59079170e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "187"
      ]
     },
     "execution_count": 68,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.argmin(np.linalg.norm(query - dataset, axis=1))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "30093128-9135-4cee-8230-02fe7136a4d2",
   "metadata": {
    "tags": []
   },
   "source": [
    "#### IVF"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "15e80847-28fb-47eb-a622-cc3b1a980176",
   "metadata": {},
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'kmeans2' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[2], line 2\u001b[0m\n\u001b[1;32m      1\u001b[0m num_part \u001b[38;5;241m=\u001b[39m \u001b[38;5;241m100\u001b[39m\n\u001b[0;32m----> 2\u001b[0m (centroids, assignments) \u001b[38;5;241m=\u001b[39m \u001b[43mkmeans2\u001b[49m(dataset, num_part, \u001b[38;5;28miter\u001b[39m\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m10000\u001b[39m)\n",
      "\u001b[0;31mNameError\u001b[0m: name 'kmeans2' is not defined"
     ]
    }
   ],
   "source": [
    "num_part = 100\n",
    "(centroids, assignments) = kmeans2(dataset, num_part, iter=10000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 71,
   "id": "9c08dfde-50af-479c-8526-c04ebd745197",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(100, 128)"
      ]
     },
     "execution_count": 71,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "centroids.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 72,
   "id": "4a4578bb-b98f-4d73-b8af-e80226066ed0",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([28, 95, 92, 79, 51, 82, 15, 53, 51, 89], dtype=int32)"
      ]
     },
     "execution_count": 72,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "assignments[:10]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 84,
   "id": "e2c70a9d-9760-4927-8761-e5eb4c3380e5",
   "metadata": {},
   "outputs": [],
   "source": [
    "index = [[] for _ in range(num_part)]\n",
    "for n, k in enumerate(assignments):\n",
    "     index[k].append(n)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 75,
   "id": "abb7097f-675b-4554-b282-7aab16dea2b6",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[14, 413]"
      ]
     },
     "execution_count": 75,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "index[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 74,
   "id": "5b0f7242-62dd-4697-942a-ba5c01b75839",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[26, 71, 104, 120, 200, 324, 540, 732, 877]"
      ]
     },
     "execution_count": 74,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "index[1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 76,
   "id": "6fdf4c34-7736-49f2-a2c3-126a14f79912",
   "metadata": {},
   "outputs": [],
   "source": [
    "cluster_id = np.argmin(np.linalg.norm(query - centroids, axis=1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 77,
   "id": "c11ea292-1871-4b30-a32d-d82edad9bc0e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "30"
      ]
     },
     "execution_count": 77,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cluster_id"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 78,
   "id": "73f29f8c-3b6f-4fe0-8662-e56837d01fb1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "9"
      ]
     },
     "execution_count": 78,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.argmin(np.linalg.norm(query - dataset[index[30]], axis=1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 102,
   "id": "86568445-8566-4ca7-b422-12aa4afbd957",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "9"
      ]
     },
     "execution_count": 102,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "index[9]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7e0fff81-840c-4e9c-a2e4-5ff8dcce756d",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 85,
   "id": "87d24129-d60a-49bf-bb61-2474e66b91b7",
   "metadata": {},
   "outputs": [],
   "source": [
    "cluster_ids = np.argsort(np.linalg.norm(query - centroids, axis=1))[: 3]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 80,
   "id": "d8a0f006-dc07-4a22-a015-fdc06a3c3026",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([30, 73, 15])"
      ]
     },
     "execution_count": 80,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cluster_ids"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 86,
   "id": "52f28985-8af6-42fa-814c-fc408ff685aa",
   "metadata": {},
   "outputs": [],
   "source": [
    "top3_index = []\n",
    "for c in cluster_ids:\n",
    "    top3_index += index[c]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 90,
   "id": "a3772600-782e-47fd-b35d-55fd6320db94",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "66"
      ]
     },
     "execution_count": 90,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.argmin(np.linalg.norm(query - dataset[top3_index], axis=1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 91,
   "id": "c3c44578-285e-483d-9d6f-5fad3ad2afa3",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "187"
      ]
     },
     "execution_count": 91,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "top3_index[66]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "67c27cf8-d214-4e72-adbf-a0294a0c9cea",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "llm",
   "language": "python",
   "name": "llm"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.19"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
